import torch.nn as nn
import torch
import torch.nn.functional as F
# 设置环境变量减少内存碎片化
import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

# 搭建unet 网络
class DoubleConv(nn.Module):  # 连续两次卷积
    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),  # 用 BN 代替 Dropout
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)


class Down(nn.Module):  # 下采样
    def __init__(self, in_channels, out_channels):
        super(Down, self).__init__()
        self.downsampling = nn.Sequential(
            nn.MaxPool2d(kernel_size=2, stride=2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.downsampling(x)


class Up(nn.Module):  # 上采样
    def __init__(self, in_channels, out_channels):
        super(Up, self).__init__()
        self.upsampling = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)  # 转置卷积
        self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.upsampling(x1)
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]
        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2])
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class OutConv(nn.Module):  # 最后一个网络的输出
    def __init__(self, in_channels, num_classes):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, num_classes, kernel_size=1)

    def forward(self, x):
        return self.conv(x)


class UNet(nn.Module):  # unet 网络
    def __init__(self, in_channels=1, num_classes=1):
        super(UNet, self).__init__()
        self.in_channels = in_channels
        self.num_classes = num_classes

        self.in_conv = DoubleConv(in_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        self.down4 = Down(512, 1024)
        self.up1 = Up(1024, 512)
        self.up2 = Up(512, 256)
        self.up3 = Up(256, 128)
        self.up4 = Up(128, 64)
        self.out_conv = OutConv(64, num_classes)

    def forward(self, x, return_features=False):
        x1 = self.in_conv(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)

        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)

        if return_features:
            return x  # 返回最后一个特征图
        else:
            return self.out_conv(x)  # 返回最终的分割结果

class PatchEmbed(nn.Module):
    def __init__(self, image_size, patch_size, in_channels, embed_dim):
        super(PatchEmbed, self).__init__()
        self.patch_size = patch_size
        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        x = self.proj(x)  # [B, C, H, W] -> [B, E, N]
        x = x.flatten(2)  # [B, E, N]
        x = x.transpose(1, 2)  # [B, N, E]
        return x

class Attention(nn.Module):
    def __init__(self, embed_dim, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0.):
        super(Attention, self).__init__()
        self.num_heads = num_heads
        head_dim = embed_dim // num_heads
        self.scale = head_dim ** -0.5
        self.qkv = nn.Linear(embed_dim, embed_dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(embed_dim, embed_dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super(Mlp, self).__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

class Block(nn.Module):
    def __init__(self, embed_dim, num_heads, mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0.):
        super(Block, self).__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn = Attention(embed_dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
        self.drop_path = nn.Dropout(drop)
        self.norm2 = nn.LayerNorm(embed_dim)
        mlp_hidden_dim = int(embed_dim * mlp_ratio)
        self.mlp = Mlp(in_features=embed_dim, hidden_features=mlp_hidden_dim, act_layer=nn.GELU, drop=drop)

    def forward(self, x):
        x = x + self.drop_path(self.attn(self.norm1(x)))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x

class VisionTransformer(nn.Module):
    def __init__(self, image_size, patch_size, in_channels, embed_dim, depth, num_heads, mlp_ratio=4., qkv_bias=True, drop_rate=0., attn_drop_rate=0.):
        super(VisionTransformer, self).__init__()
        self.patch_embed = PatchEmbed(image_size, patch_size, in_channels, embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, (image_size // patch_size) ** 2 + 1, embed_dim))
        self.pos_drop = nn.Dropout(p=drop_rate)
        self.blocks = nn.ModuleList([
            Block(embed_dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate, attn_drop=attn_drop_rate)
            for _ in range(depth)
        ])
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x):
        B, C, H, W = x.shape
        x = self.patch_embed(x)  # [B, N, E]
        N = (H // self.patch_embed.patch_size) * (W // self.patch_embed.patch_size)
        pos_embed = self.pos_embed[:, :N+1, :].detach()  # 动态调整位置嵌入的尺寸
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + pos_embed
        x = self.pos_drop(x)
        for blk in self.blocks:
            x = blk(x)
        x = self.norm(x)
        return x

# 搭建UNet_ViT网络
class UNet_ViT(nn.Module):
    def __init__(self, in_channels, num_classes, image_size, patch_size, vit_embed_dim, vit_depth, vit_heads):
        super(UNet_ViT, self).__init__()
        self.unet = UNet(in_channels, num_classes)
        self.vit = VisionTransformer(image_size=image_size, patch_size=patch_size, in_channels=in_channels, embed_dim=vit_embed_dim, depth=vit_depth, num_heads=vit_heads)
        self.fusion_conv = nn.Conv2d(vit_embed_dim + 64, 512, kernel_size=1)  # 调整输入通道数为832
        self.out_conv = nn.Conv2d(512, num_classes, kernel_size=1)  # 最终输出层
        self.image_size = image_size  # 存储image_size参数
        self.patch_size = patch_size  # 存储patch_size参数

    def forward(self, x):
        # UNet前向传播
        unet_features = self.unet(x, return_features=True)
        
        # ViT前向传播
        vit_features = self.vit(x)
        vit_features = vit_features[:, 1:, :]  # 去掉cls_token
        vit_features = vit_features.permute(0, 2, 1)  # [B, E, N]
        
        # 计算特征图的尺寸
        B, E, N = vit_features.shape
        H = (self.image_size // self.patch_size)
        W = (565 // self.patch_size)  # 使用实际的宽度
        vit_features = vit_features.view(B, E, H, W)

        # 如果需要，可以在这里添加上采样或下采样层
        if vit_features.size(2) != unet_features.size(2) or vit_features.size(3) != unet_features.size(3):
            vit_features = F.interpolate(vit_features, size=(unet_features.size(2), unet_features.size(3)), mode='bilinear', align_corners=True)

        # 特征融合
        fused_features = torch.cat((unet_features, vit_features), dim=1)
        
        # 确保融合后的特征图通道数为832
        assert fused_features.size(1) == 832, f"Expected 832 channels, but got {fused_features.size(1)} channels"

        fused_features = self.fusion_conv(fused_features)
        output = self.out_conv(fused_features)
        return output